import os
import random
import logging
import numpy as np

from datetime import datetime

from util.util_log import setup_logging
from util.util_csv import save_csv
from util.util_image import save_img, save_img_xyz, save_concat_images_horizontal_2
from util.util_ris_pattern import parse_pattern_dbw, phase_2_pattern, phase_2_pattern_xyz, \
    point_2_phase, phase_2_bit, phase_2_pattern_xyz_fft, eps, phaseBit_2_deg
from util.util_analysis_plane import set_theta_phi_loop, get_peaks, get_peak_3rd

from multi_beam.multi_beam_GA_QE import main_multi_beam_N, ga_qe_beam_N


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../../files/logs/log_multi_beam_GA_QE.log")
# 获取日志记录器并记录日志
logger = logging.getLogger("[RIS-multi-beam-GA_QE-1bit]")


# ============================================= 主函数 ====================================
# 量化选举法 - FFT加速 -- 双波束
def main_multi_scan_fft_2(path_pre,
                          theta1_begin, theta1_over, theta1_step, phi1_begin, phi1_over, phi1_step,
                          theta2_begin, theta2_over, theta2_step, phi2_begin, phi2_over, phi2_step,
                          bit_num):
    logger.info("main_multi_scan_fft_2:")
    logger.info("theta1_begin=%d, theta1_over=%d, theta1_step=%d, phi1_begin=%d, phi1_over=%d, phi1_step=%d, "
                % (theta1_begin, theta1_over, theta1_step, phi1_begin, phi1_over, phi1_step))
    logger.info("theta2_begin=%d, theta2_over=%d, theta2_step=%d, phi2_begin=%d, phi2_over=%d, phi2_step=%d, "
                % (theta2_begin, theta2_over, theta2_step, phi2_begin, phi2_over, phi2_step))
    # 轮询扫描范围, 完成两个任务: 1.找到最大的副瓣，并保存当时角度, 码阵和方向图; 2.每个角度记录方向图做动图
    theta1_loops = set_theta_phi_loop(theta1_begin, theta1_over, theta1_step)
    phi1_loops = set_theta_phi_loop(phi1_begin, phi1_over, phi1_step)
    theta2_loops = set_theta_phi_loop(theta2_begin, theta2_over, theta2_step)
    phi2_loops = set_theta_phi_loop(phi2_begin, phi2_over, phi2_step)
    i = 0
    # 收敛曲线
    list_peak_3rd = list()
    # 最大旁瓣结果
    worst_theta1 = -1
    worst_phi1 = -1
    worst_theta2 = -1
    worst_phi2 = -1
    worst_peak_3rd = -1000
    worst_phaseBit_mix = None
    worst_phaseBit_mix_deg = None
    # loop
    for theta1 in theta1_loops:
        for phi1 in phi1_loops:
            for theta2 in theta2_loops:
                for phi2 in phi2_loops:
                    logger.info("--------------------------------------------------------")
                    logger.info("(theta1, phi1)=(%d, %d), (theta2, phi2)=(%d, %d) - %d"
                                % (theta1, phi1, theta2, phi2, i))
                    i = i + 1
                    # 完成两个任务: 1.找到最大的副瓣，并保存当时角度, 码阵和方向图; 2.每个角度记录方向图做动图
                    # 计算phase1和phase2
                    phase1 = point_2_phase(theta1, phi1)
                    phase2 = point_2_phase(theta2, phi2)
                    # GA_QE
                    phaseBit1, phaseDeg1 = phase_2_bit(phase1, bit_num)
                    phaseBit2, phaseDeg2 = phase_2_bit(phase2, bit_num)
                    phaseBit_mix, best_fitness, best_fitness_history, best_individual_history = \
                        ga_qe_beam_N([phaseBit1, phaseBit2], 2, bit_num)
                    phaseBit_mix_deg = phaseBit_2_deg(phaseBit_mix, bit_num)
                    # # 计算phase_mix的方向图
                    patternBit_mix_xyz, x, y = phase_2_pattern_xyz_fft(np.deg2rad(phaseBit_mix_deg))
                    #
                    patternBit_mix_dbw = 20 * np.log10(np.abs(patternBit_mix_xyz)
                                                       / np.max(np.max(np.abs(patternBit_mix_xyz))) + eps)
                    # 1. 每个角度记录方向图做动图
                    path_img_name = "(" + str(theta1) + "," + str(phi1) + "," + str(theta2) + "," + str(phi2) + ")"
                    path_img_patternBit_xyz = path_pre + "patternBit_mix_xyz_fft/" + path_img_name + ".jpg"
                    save_img_xyz(path_img_patternBit_xyz, np.abs(patternBit_mix_xyz), x, y)
                    path_img_phaseBit = path_pre + "phaseBit_mix/" + path_img_name + ".jpg"
                    save_img(path_img_phaseBit, phaseBit_mix)
                    path_img_phase_pattern = path_pre + "phase_pattern/" + path_img_name + ".jpg"
                    save_concat_images_horizontal_2(path_img_phaseBit, path_img_patternBit_xyz, path_img_phase_pattern)
                    # 2. 找到最大的副瓣，并保存当时角度, 码阵和方向图
                    peaks = get_peaks(patternBit_mix_dbw)
                    logger.info("peaks.len=%d" % (len(peaks)))
                    if len(peaks) <= 2:
                        logger.info("ATTENTION: no peak!!!!!!!")
                        break
                    logger.info("1st: %f - (%d, %d), 2nd: %f - (%d, %d), 3rd: %f - (%d, %d)"
                                % (peaks[0][0], peaks[0][1][0], peaks[0][1][1],
                                   peaks[1][0], peaks[1][1][0], peaks[1][1][1],
                                   peaks[2][0], peaks[2][1][0], peaks[2][1][1]))
                    peak_3rd = get_peak_3rd(peaks)
                    if peak_3rd is None:
                        break
                    peak_3rd_val = peak_3rd[0]
                    peak_3rd_x, peak_3rd_y = peak_3rd[1][0], peak_3rd[1][1]
                    # list_peak_3rd.append(peak_3rd_val)
                    peak_record_len = min(6, len(peaks))
                    list_peak_3rd.append([peak_3rd_val, peak_3rd_x, peak_3rd_y, theta1, phi1, theta2, phi2,
                                            peaks[0][0], peaks[0][1][0], peaks[0][1][1],
                                            peaks[1][0], peaks[1][1][0], peaks[1][1][1],
                                            peaks[2][0], peaks[2][1][0], peaks[2][1][1],
                                            peaks[:peak_record_len]
                                            ])
                    logger.info("peak_3rd_val: %f - (%d, %d)" % (peak_3rd_val, peak_3rd_x, peak_3rd_y))
                    if worst_peak_3rd < peak_3rd_val:
                        worst_peak_3rd = peak_3rd_val
                        worst_theta1 = theta1
                        worst_phi1 = phi1
                        worst_theta2 = theta2
                        worst_phi2 = phi2
                        worst_phaseBit_mix = phaseBit_mix
                        worst_phaseBit_mix_deg = phaseBit_mix_deg
    # 保存最大旁瓣的结果
    worst_patternBit_mix = phase_2_pattern(np.deg2rad(worst_phaseBit_mix_deg))
    worst_patternBit_mix_dbw = parse_pattern_dbw(np.deg2rad(worst_phaseBit_mix_deg))
    worst_patternBit_mix_xyz, worst_patternBit_mix_xyz_x, worst_patternBit_mix_xyz_y, worst_patternBit_mix_xyz_z \
        = phase_2_pattern_xyz(np.deg2rad(worst_phaseBit_mix_deg))
    worst_peaks = get_peaks(worst_patternBit_mix_dbw)
    worst_peak_3rd = get_peak_3rd(worst_peaks)
    worst_peak_3rd_val = worst_peak_3rd[0]
    worst_peak_3rd_x, worst_peak_3rd_y = worst_peak_3rd[1][0], worst_peak_3rd[1][1]
    save_img(path_pre + "worst_phaseBit_mix.jpg", worst_phaseBit_mix)
    save_img(path_pre + "worst_patternBit_mix.jpg", worst_patternBit_mix)
    save_img(path_pre + "worst_patternBit_mix_dbw.jpg", worst_patternBit_mix_dbw)
    save_img_xyz(path_pre + "worst_patternBit_mix_xyz.jpg", np.abs(worst_patternBit_mix_xyz),
                 worst_patternBit_mix_xyz_x, worst_patternBit_mix_xyz_y)
    # 写结果
    with open(path_pre + 'worst_peak_3rd.txt', 'a') as file:
        file.write("worst: (theta1, phi1)=(%f, %f), (theta2, phi2)=(%f, %f)\n"
                   % (worst_theta1, worst_phi1, worst_theta2, worst_phi2))
        file.write("worst_peak_3rd_val=%f, worst_peak_3rd_x=%d, worst_peak_3rd_y=%d\n"
                   % (worst_peak_3rd_val, worst_peak_3rd_x, worst_peak_3rd_y))
        peak_record_len = min(6, len(worst_peaks))
        file.write("worst_peaks=%s\n" % worst_peaks[:peak_record_len])
    # 写csv
    save_csv(list_peak_3rd, path_pre + "list_peak_3rd.csv")
    save_csv(worst_phaseBit_mix, path_pre + "worst_phaseBit_mix.csv")
    save_csv(worst_patternBit_mix, path_pre + "worst_patternBit_mix.csv")
    save_csv(worst_patternBit_mix_dbw, path_pre + "worst_patternBit_mix_dbw.csv")


# 量化选举法 -- 随机角度 -- N波束
def main_multi_rand_N(path_pre, beam_num, bit_num):
    logger.info("GA_QE.main_multi_rand_N:")
    # 获取当前时间并格式化为字符串
    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    folder_name = f"{current_time}"
    folder_path = os.path.join(path_pre, folder_name) + "/"
    print(folder_path)
    # 创建文件夹
    try:
        os.makedirs(folder_path, exist_ok=True)
        logger.info(f"Created folder: {folder_path}")
    except Exception as e:
        logger.error(f"Failed to create folder: {e}")
        return
    # 生成beam_num组数，每组数由2个整数组成
    # 第一个数范围[0, 90]，第二个数范围[0, 360]
    beams = []
    for _ in range(beam_num):
        first_number = random.randint(0, 90)
        second_number = random.randint(0, 360)
        beams.append([first_number, second_number])
    # 打印生成的数
    logger.info(f"Generated {beam_num} beams: {beams}")
    # 记录生成的波束指向到.log
    with open(folder_path + 'point.txt', 'a') as file:
        for i in range(len(beams)):
            beam = beams[i]
            file.write("beam[%d]: theta0 = %d, phi0 = %d\n" % (i, beam[0], beam[1]))
    # 多波束合成
    main_multi_beam_N(beams, folder_path, bit_num)



# ======================================================= main 主方法 ===============================================
def main_multi_scan_ga_qe():
    # 波束扫描 -- 量化选举法 - FFT 加速: 主函数
    main_multi_scan_fft_2(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/2-(30-10-5)-(0,180-10-5)/",
                          20, 40, 5, -10, 10, 5, 20, 40, 5, 170, 190, 5,
                          1)
    main_multi_scan_fft_2(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/2-(30-10-5)-(0,90-10-5)/",
                          20, 40, 5, -10, 10, 5, 20, 40, 5, 80, 100, 5,
                          1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-1/", 1, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-2/", 2, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-3/", 3, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-4/", 4, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-5/", 5, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-6/", 6, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-7/", 7, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-8/", 8, 1)
    main_multi_rand_N(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/rand-16/", 16, 1)


# ======================================================= 测试方法 ===============================================
def test_scan_fft_2():
    main_multi_scan_fft_2(f"E://project/ant_sim/files/multi-beam-scan/1bit/GA_QE/2-(30-10-5)-(0,180-10-5)/",
                          40, 40, 5, -10, 10, 10, 40, 40, 5, 170, 190, 10,
                          1)

def test_rand_N():
    main_multi_rand_N(f"../files/multi-beam-scan/1bit/GA_QE/rand-5/", 5, 1)




if __name__ == '__main__':
    logger.info("1bit-RIS-multi-beam-GA_QE: Geometric partitioning method")
    # 主函数
    main_multi_scan_ga_qe()
    # 测试函数
    # test_scan_fft_2()
    # test_rand_N()